/*
 * @(#) $Header$
 *
 * Copyright (C) 2007  Forklabs Daniel Léonard
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
 */

package ca.forklabs.media.jai.opimage;

import java.awt.image.DataBuffer;
import java.awt.image.Raster;
import java.awt.image.RenderedImage;
import java.awt.image.SampleModel;
import java.lang.reflect.Array;
import java.lang.reflect.Field;
import java.util.Collection;
import javax.media.jai.CollectionImage;
import javax.media.jai.operator.DFTDataNature;
import javax.media.jai.operator.DFTScalingType;
import com.sun.media.jai.opimage.FFT;
import com.sun.media.jai.util.DataBufferUtils;
import ca.forklabs.baselib.util.Arrays;
import ca.forklabs.media.jai.RasterAdapter;
import ca.forklabs.media.jai.SimpleCollectionImage;
import ca.forklabs.media.jai.operator.DFT3DDescriptor;

/**
 * Class {@code DFT3DOpImage} is the {@link CollectionImage} for both operator
 * <em>dft3d</em> and <em>idft3d</em>.
 *
 * @author   <a href="mailto:forklabs at gmail.com?subject=ca.forklabs.media.jai.opimage.DFT3DOpImage">Daniel Léonard</a>
 * @version $Revision$
 */
public class DFT3DOpImage extends SimpleCollectionImage {

//---------------------------
// Inner classes
//---------------------------

   /**
    * Class {@code Transform} configures {@link FFT} objects.
    */
   public static abstract class Transform {

      /**
       * Constructor.
       */
      protected Transform() {
      // nothing
         }

      /**
       * Creates a newly configured {@link FFT}.
       * @param   scaling   the scaling type.
       * @return   a newly configured {@link FFT}.
       */
      public abstract FFT getDefaultTransform(DFTScalingType scaling);

      }

// there are more inner classes at the end

//---------------------------
// Class variables
//---------------------------

   /** Configures forward transforms. */
   public static final Transform FORWARD_TRANSFORM = new Transform() {
      @Override
      public FFT getDefaultTransform(DFTScalingType scaling) {
         boolean forward_transform = true;
         Integer scaling_factor = Integer.valueOf(scaling.getValue());
         int len = 2;
         FFT fft = new FFT(forward_transform, scaling_factor, len);
         return fft;
         }
      };

      /** Configures inverse transforms. */
   public static final Transform INVERSE_TRANSFORM = new Transform() {
      @Override
      public FFT getDefaultTransform(DFTScalingType scaling) {
         boolean inverse_transform = false;
         Integer scaling_factor = Integer.valueOf(scaling.getValue());
         int len = 2;
         FFT fft = new FFT(inverse_transform, scaling_factor, len);
         return fft;
         }
      };

// there are more class variables at the end

//---------------------------
// Constructors
//---------------------------

   /**
    * Constructor.
    * @param   sources   the list of sources.
    * @param   transform   the {@link FFT} configurator.
    * @param   scaling   the desired scaling.
    * @param   nature   the data nature.
    */
   public DFT3DOpImage(CollectionImage sources, Transform transform, DFTScalingType scaling, DFTDataNature nature) {
      this.transform(sources, transform, scaling, nature);
      }


//---------------------------
// Instance methods
//---------------------------

   /**
    * Performs the 3D Fourier transform.
    * @param   <T>   the type of pixel.
    * @param   sources   the source images.
    * @param   fft   the transform object.
    * @param   nature   the data nature.
    * @param   specialization   the pixel specialization.
    */
   @SuppressWarnings({ "hiding", "unchecked", "rawtypes" })
   protected <T> void transform(CollectionImage sources, FFT fft, DFTDataNature nature, PixelSpecialization<T> specialization) {
      int width, height, depth, bands, data_type;
      {
         RenderedImage image = (RenderedImage) sources.get(0);

         width = image.getWidth();
         height = image.getHeight();
         depth = sources.size();

         SampleModel sample_model = image.getSampleModel();
         bands = sample_model.getNumBands();

         data_type = sample_model.getDataType();
      }
      int len = width * height * depth;
      int image_stride = width * height;

      DataBuffer[] source_buffers = new DataBuffer[depth];
      DataBuffer[] sink_buffers = new DataBuffer[depth];
      for (int i = 0; i < depth; i++) {
         RenderedImage image = (RenderedImage) sources.get(i);
         Raster raster = image.getData();

         DataBuffer source_buffer = raster.getDataBuffer();
         source_buffers[i] = source_buffer;

         DataBuffer sink_buffer = specialization.newSinkBuffer(image_stride, bands, nature);
         sink_buffers[i] = sink_buffer;
         }

      T real_data = specialization.newWorkingArray(len);
      T imag_data = specialization.newWorkingArray(len);

   // if the source is complex, the complex data is stored as consecutive bands
      int band_stride = (specialization.isComplexSource(nature) ? 2 : 1);
      for (int b = 0; b < bands; b += band_stride) {
         for (int d = 0, offset = 0; d < depth; d++, offset += image_stride) {
            DataBuffer buffer = source_buffers[d];
            specialization.fill(real_data, imag_data, buffer, b, offset, nature);
            }

      // the row transform
         boolean should_do_row_transform = (1 != width);
         if (should_do_row_transform) {
            fft.setLength(width);
            for (int i = 0, num_rows = height * depth; i < num_rows; i++) {
               fft.setData(data_type,
                           real_data, i * width, 1,
                           imag_data, i * width, 1,
                           width);
               fft.transform();
               fft.getData(data_type,
                           real_data, i * width, 1,
                           imag_data, i * width, 1);
               }
            }

      // the column transform
         boolean should_do_col_transform = (1 != height);
         if (should_do_col_transform) {
            fft.setLength(height);
         // the start of each colum is on the first
         // row of each image, the striding is special
            for (int k = 0; k < depth; k++) {
               for (int i = 0; i < width; i++) {
                  int start = k * image_stride + i;
                  fft.setData(data_type,
                              real_data, start, width,
                              imag_data, start, width,
                              width);
                  fft.transform();
                  fft.getData(data_type,
                              real_data, start, width,
                              imag_data, start, width);
                  }
               }
            }

      // the rod transform
         boolean should_do_rod_transform = (1 != depth);
         if (should_do_rod_transform) {
            fft.setLength(depth);
            for (int i = 0, num_rods = width * height; i < num_rods; i++) {
               fft.setData(data_type,
                           real_data, i, width * height,
                           imag_data, i, width * height,
                           depth);
               fft.transform();
               fft.getData(data_type,
                           real_data, i, width * height,
                           imag_data, i, width * height);
               }
            }

      // put data into the sink buffers
         for (int d = 0, offset = 0; d < depth; d++, offset += image_stride) {
            DataBuffer buffer = sink_buffers[d];
            specialization.drain(real_data, imag_data, buffer, b, offset, nature);
            }

         }

      Collection sinks = this.getImages();
      for (int i = 0; i < depth; i++) {
         DataBuffer buffer = sink_buffers[i];
         RenderedImage sink = RasterAdapter.buildImage(buffer, width, height);
         sinks.add(sink);
         }
      }

   protected void transformFloat(CollectionImage sources, FFT fft, DFTDataNature nature) {
      PixelSpecialization<float[]> specialization = DFT3DOpImage.FLOAT_SPECIALIZATION;
      this.transform(sources, fft, nature, specialization);
      }

   protected void transformDouble(CollectionImage sources, FFT fft, DFTDataNature nature) {
      PixelSpecialization<double[]> specialization = DFT3DOpImage.DOUBLE_SPECIALIZATION;
      this.transform(sources, fft, nature, specialization);
      }

   protected void transform(CollectionImage sources, Transform transform, DFTScalingType scaling, DFTDataNature nature) {
      FFT fft = transform.getDefaultTransform(scaling);

      RenderedImage image = (RenderedImage) sources.get(0);
      SampleModel sample_model = image.getSampleModel();
      int data_type = sample_model.getDataType();
      switch (data_type) {
         case DataBuffer.TYPE_FLOAT:
            this.transformFloat(sources, fft, nature);
            break;
         case DataBuffer.TYPE_DOUBLE:
            this.transformDouble(sources, fft, nature);
            break;
         default:
            String message = this.getBadDataType(data_type);
            throw new IllegalStateException(message);
         }
      }

   @SuppressWarnings("boxing")
   protected String getBadDataType(int type) {
      String key = Resources.DFT3D_BAD_DATA_TYPE;
      String message = Resources.getLocalizedString(key, type);
      return message;
      }


//---------------------------
// Inner classes
//---------------------------

   /**
    * Class {@link PixelSpecialization} specializes the work on the type of
    * pixels.
    * @param   <T>   the type of pixel.
    */
   protected static abstract class PixelSpecialization<T> {

      /**
       * Fills both the real and imaginary work arrays with all the data of the
       * desired band in row order.
       * @param   real   the real work array.
       * @param   imag   the imaginary work array.
       * @param   buffer   the data buffer containing the data.
       * @param   band   the desired band.
       * @param   offset   the starting offset in both work arrays.
       * @param   nature   the data nature.
       */
      public abstract void fill(T real, T imag, DataBuffer buffer, int band, int offset, DFTDataNature nature);

      /**
       * Drains both the real and imaginary work arrays into the buffer.
       * @param   real   the real work array.
       * @param   imag   the imaginary work array.
       * @param   buffer   the data buffer containing the data.
       * @param   band   the desired band.
       * @param   offset   the starting offset in both work arrays.
       * @param   nature   the data nature.
       */
      public abstract void drain(T real, T imag, DataBuffer buffer, int band, int offset, DFTDataNature nature);

      /**
       * Creates a new work array of the desired length.
       * @param   len   the length.
       * @return   a new array of the desired length.
       */
      public abstract T newWorkingArray(int len);

      /**
       * Creates a {@link DataBuffer} of the specialized type big enough to hold
       * the data.
       * @param   size   the size of each band.
       * @param   bands   the number of bands.
       * @param   nature   the data nature, complex images must have twice the
       *                   number of desired bands.
       * @return   a new data buffer.
       */
      public abstract DataBuffer newSinkBuffer(int size, int bands, DFTDataNature nature);

      /**
       * Gets the raw array inside a {@link DataBuffer}.
       * @param   name   the name of the field.
       * @param   buffer   the object to search into.
       * @return   the value of the field, the raw array inside the data buffer.
       */
      protected Object getBankArray(String name, DataBuffer buffer) {
         Object value;
         try {
            Class<?> clazz = buffer.getClass();

            Field field = clazz.getDeclaredField(name);
            field.setAccessible(true);

            value = field.get(buffer);
            }
         catch (Exception e) {
            throw new IllegalStateException(e);
            }
         return value;
         }

      /**
       * Determines if the nature is of type <em>REAL_TO_COMPLEX</em>.
       * @param    nature   the data nature.
       * @return   {@code true} if the nature is of type
       *           <em>REAL_TO_COMPLEX</em>, {@code false} otherwise.
       */
      protected boolean isRealToComplex(DFTDataNature nature) {
         boolean is_real_to_complex = (nature == DFT3DDescriptor.REAL_TO_COMPLEX);
         return is_real_to_complex;
         }

      /**
       * Determines if the nature is of type <em>COMPLEX_TO_COMPLEX</em>.
       * @param    nature   the data nature.
       * @return   {@code true} if the nature is of type
       *           <em>COMPLEX_TO_COMPLEX</em>, {@code false} otherwise.
       */
      protected boolean isComplexToComplex(DFTDataNature nature) {
         boolean is_complex_to_complex = (nature == DFT3DDescriptor.COMPLEX_TO_COMPLEX);
         return is_complex_to_complex;
         }

      /**
       * Determines if the nature is of type <em>COMPLEX_TO_REAL</em>.
       * @param    nature   the data nature.
       * @return   {@code true} if the nature is of type
       *           <em>COMPLEX_TO_REAL</em>, {@code false} otherwise.
       */
      protected boolean isComplexToReal(DFTDataNature nature) {
         boolean is_complex_to_real = (nature == DFT3DDescriptor.COMPLEX_TO_REAL);
         return is_complex_to_real;
         }

      /**
       * Determines from the nature if the source is of real type.
       * @param   nature   the data nature.
       * @return   {@code true} of the source if or real type, {@code false}
       *           otherwise.
       */
      public boolean isRealSource(DFTDataNature nature) {
         boolean is_real_to_complex = this.isRealToComplex(nature);
         boolean is_real_source = is_real_to_complex;
         return is_real_source;
         }

      /**
       * Determines from the nature if the source is of complex type.
       * @param   nature   the data nature.
       * @return   {@code true} of the source if or complex type, {@code false}
       *           otherwise.
       */
      public boolean isComplexSource(DFTDataNature nature) {
         boolean is_complex_to_real = this.isComplexToReal(nature);
         boolean is_complex_to_complex = this.isComplexToComplex(nature);
         boolean is_complex_source = is_complex_to_real || is_complex_to_complex;
         return is_complex_source;
         }

      /**
       * Determines from the nature if the destination is of complex type.
       * @param   nature   the data nature.
       * @return   {@code true} of the destination if or complex type,
       *           {@code false} otherwise.
       */
      public boolean isComplexDestination(DFTDataNature nature) {
         boolean is_real_to_complex = this.isRealToComplex(nature);
         boolean is_complex_to_complex = this.isComplexToComplex(nature);
         boolean is_complex_destination = is_real_to_complex || is_complex_to_complex;
         return is_complex_destination;
         }

      /**
       * Calculates the number of band of the sink buffer.
       * @param   bands   the number of bands in the sources.
       * @param   nature   the data nature, complex buffer have twice the
       *          desired number of bands.
       * @return   the adjusted number of bands.
       */
      protected int calculateSinkBanks(int bands, DFTDataNature nature) {
      // the number of bands is from the source
      // image, it is used to calculate the number
      // of band of the destination image
         boolean is_real_to_complex = this.isRealToComplex(nature);
         boolean is_complex_to_complex = this.isComplexToComplex(nature);
         boolean is_complex_to_real = this.isComplexToReal(nature);

         int banks;
         if (is_real_to_complex) {
            banks = bands * 2;
            }
         else if (is_complex_to_complex) {
            banks = bands;
            }
         else if (is_complex_to_real) {
            banks = bands / 2;
            }
         else {
            String message = this.getUnknownDataNature(nature);
            throw new IllegalStateException(message);
            }
         return banks;
         }

      /**
       * Gets the error message saying that the data nature is unknown.
       * @param   nature   the data nature.
       * @return   the error message.
       */
      protected String getUnknownDataNature(DFTDataNature nature) {
         String key = Resources.UNKNOWN_DATA_NATURE;
         String message = Resources.getLocalizedString(key, nature);
         return message;
         }

      }

   /** Specialization for float pixels. */
   protected static final PixelSpecialization<float[]> FLOAT_SPECIALIZATION = new PixelSpecialization<float[]>() {

      @Override
      @SuppressWarnings("nls")
      public void drain(float[] real, float[] imag, DataBuffer buffer, int band, int offset, DFTDataNature nature) {
      // there is no method that provide bulk access to the
      // data inside the data buffer, reflection is required
         Object value = this.getBankArray("bankdata", buffer);

      // alway drain the real part
         float[] sink = (float[]) Array.get(value, band);
         int len = sink.length;
         System.arraycopy(real, offset, sink, 0, len);

      // maybe sink the imaginary part
         boolean is_complex_destination = this.isComplexDestination(nature);
         if (is_complex_destination) {
            sink = (float[]) Array.get(value, band + 1);
            len = sink.length;
            System.arraycopy(imag, offset, sink, 0, len);
            }
         }

      @Override
      public void fill(float[] real, float[] imag, DataBuffer buffer, int band, int offset, DFTDataNature nature) {
      // there is alway a real part
         float[] data = DataBufferUtils.getDataFloat(buffer, band);
         int len = data.length;
         System.arraycopy(data, 0, real, offset, len);

      // maybe there is an imaginary part ...
         boolean is_complex_source = this.isComplexSource(nature);
         boolean is_real_source = this.isRealSource(nature);
         if (is_complex_source) {
            data = DataBufferUtils.getDataFloat(buffer, band + 1);
            len = data.length;
            System.arraycopy(data, 0, imag, offset, len);
            }
      // ... but if not, fill with zeros
         else if (is_real_source) {
            Arrays.memset(imag, 0.0f, offset, len);
            }
         else {
            String message = this.getUnknownDataNature(nature);
            throw new IllegalStateException(message);
            }
         }

      @Override
      public float[] newWorkingArray(int len) {
         float[] array = new float[len];
         return array;
         }

      @Override
      public DataBuffer newSinkBuffer(int size, int bands, DFTDataNature nature) {
         int banks = this.calculateSinkBanks(bands, nature);
         DataBuffer buffer = DataBufferUtils.createDataBufferFloat(size, banks);
         return buffer;
         }

      };

   /** Specialization for double pixels. */
   protected static final PixelSpecialization<double[]> DOUBLE_SPECIALIZATION = new PixelSpecialization<double[]>() {

      @Override
      @SuppressWarnings("nls")
      public void drain(double[] real, double[] imag, DataBuffer buffer, int band, int offset, DFTDataNature nature) {
         // there is no method that provide bulk access to the
         // data inside the data buffer, reflection is required
            Object value = this.getBankArray("bankdata", buffer);

         // alway drain the real part
            double[] sink = (double[]) Array.get(value, band);
            int len = sink.length;
            System.arraycopy(real, offset, sink, 0, len);

         // maybe sink the imaginary part
            boolean is_complex_destination = this.isComplexDestination(nature);
            if (is_complex_destination) {
               sink = (double[]) Array.get(value, band + 1);
               len = sink.length;
               System.arraycopy(imag, offset, sink, 0, len);
               }
            }

      @Override
      public void fill(double[] real, double[] imag, DataBuffer buffer, int band, int offset, DFTDataNature nature) {
      // there is alway a real part
         double[] data = DataBufferUtils.getDataDouble(buffer, band);
         int len = data.length;
         System.arraycopy(data, 0, real, offset, len);

      // maybe there is an imaginary part ...
         boolean is_complex_source = this.isComplexSource(nature);
         boolean is_real_source = this.isRealSource(nature);
         if (is_complex_source) {
            data = DataBufferUtils.getDataDouble(buffer, band + 1);
            len = data.length;
            System.arraycopy(data, 0, imag, offset, len);
            }
      // ... but if not, fill with zeros
         else if (is_real_source) {
            Arrays.memset(imag, 0.0, offset, len);
            }
         else {
            String message = this.getUnknownDataNature(nature);
            throw new IllegalStateException(message);
            }
         }

      @Override
      public double[] newWorkingArray(int len) {
         double[] array = new double[len];
         return array;
         }

      @Override
      public DataBuffer newSinkBuffer(int size, int bands, DFTDataNature nature) {
         int banks = this.calculateSinkBanks(bands, nature);
         DataBuffer buffer = DataBufferUtils.createDataBufferDouble(size, banks);
         return buffer;
         }

      };

   }
